# Goal: Estimate 2nd stage model allowing coefficients to vary by interview year ('N.draws' times, each time using a different sample from types' distribution)
# Dependencies: "datamodel.Rdata", "samples_stage1.Rdata"

setwd("~/Replication Package")

library(rjags)
library(coda)
load.module("glm")

library(foreach)
library(doMC)  
registerDoMC(4)  

RNGkind("L'Ecuyer-CMRG")

set.seed(2000)

load("datamodel.Rdata")

load("samples_stage1.Rdata")

conv.type.samples <- rbind(samples.stage1[[1]][ , substr(colnames(samples.stage1[[1]]), 1, 5) == "conv."], samples.stage1[[2]][ , substr(colnames(samples.stage1[[2]]), 1, 5) == "conv."], samples.stage1[[3]][ , substr(colnames(samples.stage1[[3]]), 1, 5) == "conv."])

unconv.type.samples <- rbind(samples.stage1[[1]][ , substr(colnames(samples.stage1[[1]]), 1, 5) == "uncon"], samples.stage1[[2]][ , substr(colnames(samples.stage1[[2]]), 1, 5) == "uncon"], samples.stage1[[3]][ , substr(colnames(samples.stage1[[3]]), 1, 5) == "uncon"])

N.draws <- 30 # NUMBER OF DRAWS FROM TYPE DISTRIBUTION

tosample <- matrix(sample(1:3000, N.draws * dim(conv.type.samples)[2], replace = T), ncol = dim(conv.type.samples)[2])

conv.type.smallsamples <- matrix(NA, ncol = dim(conv.type.samples)[2], nrow = N.draws)
unconv.type.smallsamples <- matrix(NA, ncol = dim(unconv.type.samples)[2], nrow = N.draws)

for (i in 1:dim(conv.type.samples)[2]) {
  conv.type.smallsamples[,i] <- conv.type.samples[tosample[,i],i] - 1
  unconv.type.smallsamples[,i] <- unconv.type.samples[tosample[,i],i] - 1
}

##---------- BUGS/JAGS code 2nd stage model ----------##

modelString = "

model{

  for (i in 1:n) { 

    typeComb[i] ~ dcat(p[i, 1:ncells])

    for (c in 1:ncells) {

      mu[i, c] <- beta[c, year[i], 1] +
      beta[c, year[i], 2] * education[i] +
      beta[c, year[i], 3] * age[i] +
      beta[c, year[i], 4] * female[i] +
      beta[c, year[i], 5] * nonwhite[i] +
      beta[c, year[i], 6] * income.proxy[i] +
      beta[c, year[i], 7] * ideology[i] +
      beta[c, year[i], 8] * natecon[i] +
      beta[c, year[i], 9] * persfin[i] +
      beta[c, year[i], 10] * victim[i] +
      beta[c, year[i], 11] * corrupt[i] +
      beta[c, year[i], 12] * capital[i] +
      beta[c, year[i], 13] * bigcity[i]

      emu[i, c] <- exp(mu[i, c])

      p[i, c] <- emu[i, c] / sum(emu[i, 1:ncells])

    }

  }

  for (year in 1:nyear) {
    for (k in 1:nvar) {
      beta[1, year, k] <- 0
    }

    for (c in 2:ncells) {
      beta[c, year, 1:nvar] ~ dmnorm(mu.beta[c, 1:nvar], Tau.beta[c, 1:nvar, 1:nvar])
    }
  }

  for (k in 1:nvar) {
    mu.beta[1, k] <- 0
  }

  Sigma.beta[1, 1:nvar, 1:nvar] <- B0.null

  for (c in 2:ncells) {
    mu.beta[c, 1:nvar] ~ dmnorm(b0, B0)
    Tau.beta[c, 1:nvar, 1:nvar] ~ dwish(B0, nvar)
    Sigma.beta[c, 1:nvar, 1:nvar] <- inverse(Tau.beta[c, 1:nvar, 1:nvar])
  }

}

"
writeLines(modelString, con = "stage2_reff.bug")

##---------- ESTIMATE 2-STAGE MODEL ----------##

reffs.samples <- foreach(j = 1:N.draws) %dopar% {
  
  conv.type <- conv.type.smallsamples[j,]
  unconv.type <- unconv.type.smallsamples[j,]
  
  typeComb <- NULL
  typeComb[conv.type == 0 & unconv.type == 0] <- 1
  typeComb[conv.type == 0 & unconv.type == 1] <- 2
  typeComb[conv.type == 1 & unconv.type == 0] <- 3
  typeComb[conv.type == 1 & unconv.type == 1] <- 4
  
  ncells <- 4
  
  nvar <- 13
  
  nyear <- 3
  
  b0 <- rep(0, nvar)
  B0 <- diag(nvar) * 0.01
  
  B0.null <- diag(nvar) * 0
  
  datamodel2 <- data.frame("year" = datamodel$year.cat, "typeComb" = typeComb, "education" = datamodel$education, "age" = datamodel$age, "female" = datamodel$female, "nonwhite" = datamodel$nonwhite, "income.proxy" = datamodel$income.proxy, "ideology" = datamodel$ideology, "natecon" = datamodel$natecon, "persfin" = datamodel$persfin, "victim" = datamodel$victim, "corrupt" = datamodel$corrupt, "capital" = datamodel$capital, "bigcity" = datamodel$bigcity)
  
  # Standardize covariates
  datamodel2[ , 3:14] <- scale(datamodel2[ , 3:14])
  
  n <- dim(datamodel2)[1]
  
  data2.jags  <- list("n" = n, "ncells" = ncells, "nvar" = nvar, "nyear" = nyear, "year" = datamodel2$year, "b0"= b0, "B0"= B0, "B0.null"= B0.null, "typeComb"=  datamodel2$typeComb, "education" = datamodel2$education, "age" = datamodel2$age, "female" = datamodel2$female, "nonwhite" = datamodel2$nonwhite, "income.proxy" = datamodel2$income.proxy, "ideology" = datamodel2$ideology, "natecon" = datamodel2$natecon, "persfin" = datamodel2$persfin, "victim" = datamodel2$victim, "corrupt" = datamodel2$corrupt, "capital" = datamodel2$capital, "bigcity" = datamodel2$bigcity)
  
  parameters2.jags<-c("beta", "mu.beta", "Sigma.beta")
  
  inits.jags <- list(list(.RNG.seed = sample(1:1000,1), .RNG.name = "base::Mersenne-Twister"))
  
  model2.jags <- jags.model("stage2_reff.bug", data = data2.jags, inits = inits.jags, n.chains = 1, n.adapt = 50000)
  
  mcmc.reffs.samples <- coda.samples(model2.jags, variable.names = parameters2.jags, n.iter = 50000, thin = 50)
  
  mcmc.reffs.samples[[1]]
  
}

save(reffs.samples, file = "reffs_samples.Rdata")
